import pandas as pd
import os
import argparse
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
import requests
from dotenv import load_dotenv

# 加载 .env 文件
load_dotenv()

# 设置代理（如需）
os.environ["HTTP_PROXY"] = "http://127.0.0.1:7890"
os.environ["HTTPS_PROXY"] = "http://127.0.0.1:7890"

# YouTube API 配置
YT_API_KEY = os.getenv("YT_API_KEY")
VIDEO_ID = os.getenv("TEST_VIDEO_ID", "BXCpJRiQAWs")  # 可选设置测试视频ID

# 创建服务对象
youtube = build('youtube', 'v3', developerKey=YT_API_KEY)


def get_comments(video_id):
    """
    获取指定视频的所有评论及其子评论
    :param video_id: 视频ID
    :return: 评论列表 [[author, date, likes, comment_text, parent_id], ...]
             如果是子评论，则 parent_id 为对应顶级评论的 ID；否则为 None
    """
    comments = []
    next_page_token = None

    while True:
        try:
            response = youtube.commentThreads().list(
                part="snippet,replies",
                videoId=video_id,
                maxResults=100,
                pageToken=next_page_token,
                textFormat="plainText"
            ).execute()
        except HttpError as e:
            print("请求失败:", e)
            break

        for item in response['items']:
            # 获取顶级评论
            top_comment = item['snippet']['topLevelComment']['snippet']
            top_comment_id = item['snippet']['topLevelComment']['id']
            comments.append([
                top_comment['authorDisplayName'],
                top_comment['publishedAt'],
                top_comment['likeCount'],
                top_comment['textOriginal'],
                None  # 表示这是顶级评论
            ])

            # 如果有子评论，获取它们
            if 'replies' in item:
                for reply in item['replies']['comments']:
                    reply_snippet = reply['snippet']
                    comments.append([
                        reply_snippet['authorDisplayName'],
                        reply_snippet['publishedAt'],
                        reply_snippet['likeCount'],
                        reply_snippet['textOriginal'],
                        top_comment_id  # 标记该评论为子评论，并记录父级 ID
                    ])

        next_page_token = response.get('nextPageToken')
        if not next_page_token:
            break

    return comments


class CommentAnalyzer:
    def __init__(self, model_name="deepseek"):
        self.model_name = model_name.lower()
        self.headers = {}

        if self.model_name == "deepseek":
            self.api_url = "https://api.deepseek.com/chat/completions"
            self.headers['Authorization'] = f'Bearer {os.getenv("DEEPSEEK_API_KEY")}'

        elif self.model_name == "r1":
            self.api_url = "https://api.r1.ai/v1/chat/completions"
            self.headers['Authorization'] = f'Bearer {os.getenv("R1_API_KEY")}'

        else:
            raise ValueError(f"❌ 不支持的模型类型: {model_name}")

        self.headers['Content-Type'] = 'application/json'

    def analyze_comment(self, comment_text):
        """
        调用指定模型接口进行标签分析
        """
        if not comment_text or not isinstance(comment_text, str) or not comment_text.strip():
            return []

        data = {
            "messages": [
                {"role": "user", "content": f"请为以下评论生成标签（最多3个，逗号分隔）：{comment_text}"}
            ]
        }

        if self.model_name == "deepseek":
            data["model"] = "deepseek-chat"
        elif self.model_name == "r1":
            data["model"] = "r1-mini"

        try:
            response = requests.post(self.api_url, headers=self.headers, json=data, timeout=30)
            result = response.json()

            if 'choices' in result and len(result['choices']) > 0:
                content = result['choices'][0]['message']['content'].strip()
                return [tag.strip() for tag in content.split(',') if tag.strip()]
            else:
                print("⚠️ 返回结果格式异常")
                return []

        except Exception as e:
            print(f"❌ 调用 {self.model_name} 失败: {str(e)}")
            return []


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="YouTube评论采集与分析测试脚本")
    parser.add_argument("--model", choices=["deepseek", "r1"], default="deepseek", help="使用的AI模型 (默认: deepseek)")
    parser.add_argument("--video-id", help="要抓取评论的视频ID")
    args = parser.parse_args()

    VIDEO_ID = args.video_id or VIDEO_ID

    # 获取并导出评论
    print(f"🔄 正在获取视频 {VIDEO_ID} 的评论...")
    comments_data = get_comments(VIDEO_ID)

    print(f"🧠 使用 {args.model.upper()} 模型进行标签分析...")
    analyzer = CommentAnalyzer(model_name=args.model)

    analyzed_data = []
    for idx, comment_row in enumerate(comments_data):
        comment_text = comment_row[-2]  # 获取评论内容字段（倒数第二列）
        if not comment_text or not isinstance(comment_text, str):
            print(f"⚠️ 第 {idx + 1} 条评论内容为空或无效，跳过分析")
            analyzed_data.append(comment_row + [""])
            continue

        tags = analyzer.analyze_comment(comment_text)
        analyzed_data.append(comment_row + [", ".join(t for t in tags)])
        print(f"✅ 已处理第 {idx + 1} 条评论")

    # 添加列名
    df = pd.DataFrame(analyzed_data, columns=['Author', 'Date', 'Likes', 'Comment', 'ParentID', 'Tags'])

    output_file = f"youtube_comments_with_{args.model}_tags.xlsx"
    df.to_excel(output_file, index=False, engine='openpyxl')
    print(f"✅ 导出成功：{output_file}")